import json
import os
import numpy as np
import vtk
from vedo import *

def GetVTKTransformationMatrix(matrix=[[-180, 180], [-180, 180], [-180, 180],
                                       [-10, 10], [-10, 10], [-10, 10],
                                       [0.8, 1.2], [0.8, 1.2], [0.8, 1.2]]):

    """
    get transformation matrix
    :param matrix: rotate_XYZ, translate_XYZ, scale_XYZ
    :return: vtkMatrix4x4
    """

    [rotate_X, rotate_Y, rotate_Z, translate_X, translate_Y, translate_Z, scale_X, scale_Y, scale_Z] = matrix

    Trans = vtk.vtkTransform()

    rx_flag = np.random.randint(0, 2)
    ry_flag = np.random.randint(0, 2)
    rz_flag = np.random.randint(0, 2)
    if rx_flag == 1:
        Trans.RotateX(np.random.uniform(rotate_X[0], rotate_X[1]))
    if ry_flag == 1:
        Trans.RotateY(np.random.uniform(rotate_Y[0], rotate_Y[1]))
    if rz_flag == 1:
        Trans.RotateZ(np.random.uniform(rotate_Z[0], rotate_Z[1]))

    trans_flag = np.random.randint(0, 2)
    if trans_flag == 1:
        Trans.Translate([np.random.uniform(translate_X[0], translate_X[1]),
                         np.random.uniform(translate_Y[0], translate_Y[1]),
                         np.random.uniform(translate_Z[0], translate_Z[1])])

    scale_flag = np.random.randint(0, 2)
    if scale_flag == 1:
        Trans.Scale([np.random.uniform(scale_X[0], scale_X[1]),
                     np.random.uniform(scale_Y[0], scale_Y[1]),
                     np.random.uniform(scale_Z[0], scale_Z[1])])

    matrix = Trans.GetMatrix()
    return matrix

if __name__ == '__main__':

    with open('config.json', 'r') as f:
        config = json.load(f)

    mod = config["mod"]
    num_sample = config["data"]["num_data"]
    dataset_path = config["data"]["dataset_path"]
    data_path = dataset_path + config["data"]["data_path"] + mod
    output_save_path = dataset_path + config["data"]["augmentation_data_path"] + mod

    if not os.path.exists(output_save_path):
        os.mkdir(output_save_path)

    sample_list = list(range(1, num_sample + 1))
    num_augmentations = config["num_augmentations"]

    for i_sample in sample_list:
        file_name = '{0}{1}.ply'.format(mod, i_sample)
        if not os.path.exists(os.path.join(data_path, file_name)):
            continue

        for i_aug in range(num_augmentations):

            file_name = '{0}{1}.ply'.format(mod, i_sample)
            output_file_name = 'A{0}_{1}{2}.ply'.format(i_aug, mod, i_sample)
            vtk_matrix = GetVTKTransformationMatrix()

            mesh = load(os.path.join(data_path, file_name))
            mesh.applyTransform(vtk_matrix)
            io.write(mesh, os.path.join(output_save_path, output_file_name))